package org.atlantis.datasink.utils;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public final class SubstepAwait {

	private CountDownLatch latch;
	private long timeout;

	public static SubstepAwait create(int count, long timeout) {
		return new SubstepAwait(count, timeout);
	}

	public static SubstepAwait create(int count) {
		return create(count, -1);
	}

	private SubstepAwait(int count, long timeout) {
		this.latch = new CountDownLatch(count);
		this.timeout = timeout;
	}

	public void exec(Runnable task) {
		Runner runner = new Runner();
		runner.task = task;
		new Thread(runner).start();
	}

	public void waitFinish() {
		try {
			if (timeout < 0) {
				latch.await();
				return;
			}
			if (!latch.await(timeout, TimeUnit.MILLISECONDS)) throw new IllegalStateException(String.format(
					"Some steps timeout(%d seconds).", TimeUnit.MILLISECONDS.toSeconds(timeout)));
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
		}
	}

	private class Runner implements Runnable {
		Runnable task;

		@Override
		public void run() {
			try {
				task.run();
			} finally {
				latch.countDown();
			}

		}
	}
}
